# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/8/25 上午10:32
# @Author  : shutao
# @FileName: train.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import argparse
import os
import pprint
import random
import time
import warnings
from loguru import logger

import paddle
from tabulate import tabulate

from modules.core import Trainer, launch
from modules.exp import get_exp
from modules.utils import PRAYER


logger.add('logs/{}.log'.format(int(time.time())),
           level='DEBUG',
           format='{time:YYYY-MM-DD HH:mm:ss}-{level}-{file}-{line}-{message}',
           rotation='100 MB',
           encoding="utf-8")


def make_parser():
    parser = argparse.ArgumentParser("Resnet train parser")
    parser.add_argument("-expn", "--experiment-name", type=str, default=None)
    parser.add_argument("-n", "--name", type=str, default="stgcn_s", help="model name")  # todo: 修改此处以更改不同实验

    parser.add_argument("-b", "--batch-size", type=int, default=8, help="batch size")
    parser.add_argument(
        "-f",
        "--exp_file",
        default=None,
        type=str,
        help="plz input your expriment description file",
    )
    parser.add_argument(
        "--resume", default=False, action="store_true", help="resume training"
    )
    parser.add_argument(
        "-e",
        "--start_epoch",
        default=None,
        type=int,
       help="resume training start epoch",
    )
    parser.add_argument(
        "--fp16",
        dest="fp16",
        default=False,
        action="store_true",
        help="Adopting mix precision training.",
    )
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    parser.add_argument(
        "--model_params",
        default=None,
        type=str
    )
    parser.add_argument(
        "--opt_params",
        default=None,
        type=str
    )
    return parser


@logger.catch
def main(exp, args):
    if exp.seed is not None:
        random.seed(exp.seed)
        paddle.seed(exp.seed)
        warnings.warn(
            "You have chosen to seed training. This will turn on the CUDNN deterministic setting, "
            "which can slow down your training considerably! You may see unexpected behavior "
            "when restarting from checkpoints."
        )

    # Don't remove this line of code, it will pray for your project!!!
    PRAYER()

    table_header = ["keys", "values"]
    exp_table = [(str(k), pprint.pformat(v)) for k, v in vars(args).items() if not k.startswith("_")]
    ars = tabulate(exp_table, headers=table_header, tablefmt="fancy_grid")
    logger.info("args value:\n{}".format(ars))

    trainer = Trainer(exp, args)
    trainer.train()


if __name__ == '__main__':
    args = make_parser().parse_args()
    exp = get_exp(args.exp_file, args.name)
    exp.merge(args.opts)

    if not args.experiment_name:
        args.experiment_name = exp.exp_name

    launch(main, args=(exp, args))
